/*
 * Copyright 2016-2020 dromara.org.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.dromara.mendmix.gateway.filter;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.Comparator;
import java.util.List;
import java.util.stream.Collectors;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

import org.apache.commons.lang3.StringUtils;
import org.dromara.mendmix.common.CustomRequestHeaders;
import org.dromara.mendmix.common.GlobalConstants;
import org.dromara.mendmix.common.util.JsonUtils;
import org.dromara.mendmix.gateway.GatewayConstants;
import org.dromara.mendmix.gateway.helper.HttpRequestHelper;
import org.dromara.mendmix.gateway.helper.RequestContextHelper;
import org.dromara.mendmix.gateway.model.BizSystemModule;
import org.dromara.mendmix.logging.reqlog.ApiRequestLog;
import org.dromara.mendmix.logging.reqlog.RequestLogCollector;
import org.dromara.mendmix.spring.InstanceFactory;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.util.FileCopyUtils;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.server.ServerWebExchange;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * 
 * @description <br>
 * @author <a href="mailto:vakinge@gmail.com">jiangwei</a>
 * @date 2022年4月7日
 */
public class RewriteBodyServerHttpResponse extends ServerHttpResponseDecorator {


	private static Logger logger = LoggerFactory.getLogger("org.dromara.mendmix");
	
	private static final String ATTACHMENT_FILENAME = "attachment;filename=";
	private static final String GZIP_ENCODE = "gzip";

	//private static WebClient webClient = InstanceFactory.getInstance(WebClient.class);
	private static List<HttpMessageReader<?>> messageReaders;
	 private static List<SubFilterHandler> subFilterHandlers;
	
	private ServerWebExchange exchange;
	private BizSystemModule module;
	private boolean ignoreRewriteOnSuccess;
	
	private String bodyString;
	
	private static List<HttpMessageReader<?>> getMessageReaders() {
		if(messageReaders == null) {
			//messageReaders = HandlerStrategies.withDefaults().messageReaders();
			messageReaders = InstanceFactory.getInstance(ServerCodecConfigurer.class).getReaders();
		}
		return messageReaders;
	}
	
	public static void setHandlers(List<SubFilterHandler> handlers) {
		if (handlers.size() > 1) {
			handlers = handlers.stream().sorted(Comparator.comparing(SubFilterHandler::ordered)).collect(Collectors.toList());
		}
		RewriteBodyServerHttpResponse.subFilterHandlers = handlers;
	}

	public RewriteBodyServerHttpResponse(ServerWebExchange exchange, BizSystemModule module) {
		super(exchange.getResponse());
		this.exchange = exchange;
		this.module = module;
	}

	public String getBodyString() {
		return bodyString;
	}

	@Override
	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
		if(exchange.getAttributes().containsKey(GatewayConstants.CONTEXT_IGNORE_FILTER)) {
			return super.writeWith(body);
		}
		final boolean debugMode = RequestContextHelper.isDebugMode(exchange); 
		final HttpHeaders headers = exchange.getResponse().getHeaders();
    	//
    	if(headers.containsKey(CustomRequestHeaders.HEADER_ACCEL_REDIRECT)) {
    		String proxyUrl = headers.getFirst(CustomRequestHeaders.HEADER_ACCEL_REDIRECT);
    		String contentDispostion = headers.getFirst(HttpHeaders.CONTENT_DISPOSITION);
    		if(!proxyUrl.startsWith(GatewayConstants.HTTP_SCHAME)) {
    			String[] parts = StringUtils.splitByWholeSeparator(proxyUrl, "dest=");
    			String baseUrl = parts[1];
    			String path = parts[0].substring(0,parts[0].length() - 1);
    			if(baseUrl.contains(GlobalConstants.PARAM_SEPARATOR)) {
    				baseUrl = StringUtils.split(baseUrl, GlobalConstants.PARAM_SEPARATOR, 2)[0];
    			}
    			proxyUrl = baseUrl + path;
    			if(StringUtils.isBlank(contentDispostion)) {
    				String fileName = path.substring(path.lastIndexOf(GlobalConstants.PATH_SEPARATOR) + 1);
    				if(fileName.contains(GlobalConstants.QUESTION_MASK)) {
    					fileName = StringUtils.split(fileName, GlobalConstants.QUESTION_MASK, 2)[0];
    				}
    				contentDispostion = ATTACHMENT_FILENAME + fileName;
    			}
    		}else if(StringUtils.isBlank(contentDispostion)) {
    			String fileName = URI.create(proxyUrl).getPath();
    			if(fileName.contains(GlobalConstants.PATH_SEPARATOR)) {
    				fileName = fileName.substring(fileName.lastIndexOf(GlobalConstants.PATH_SEPARATOR) + 1);
    			}
    			contentDispostion = ATTACHMENT_FILENAME + fileName;
    		}
    		Mono<DataBuffer> mono = HttpRequestHelper.proxyRequest(exchange.getResponse(), proxyUrl);
    		headers.add(HttpHeaders.CONTENT_DISPOSITION, contentDispostion);
    		return super.writeWith(mono);
//   		    Flux<DataBuffer> imageFlux = webClient.get()
//   	                .uri(proxyUrl)
//   	                .accept(MediaType.ALL) 
//   	                .retrieve()
//   	                .bodyToFlux(DataBuffer.class);
//   		   return super.writeWith(imageFlux);
    	}
    	
    	if (headers.containsKey(HttpHeaders.CONTENT_DISPOSITION) || headers.containsKey(CustomRequestHeaders.HEADER_HIT_CACHE_KEY)) {
    		return super.writeWith(body);
		}
    	//
		final long contentLength = headers.getContentLength();
		final boolean withTransferEncodingHeader = headers.containsKey(HttpHeaders.TRANSFER_ENCODING);
		boolean isNullBody = !withTransferEncodingHeader && contentLength == 0;
		if(isNullBody) {
			return super.writeWith(body);
		}

		final boolean isGzip = GZIP_ENCODE.equalsIgnoreCase(headers.getFirst(HttpHeaders.CONTENT_ENCODING));
        if(debugMode) {
        	logger.info("handleReadResponseBody begin -> isGzip:{}",isGzip);
        }

        ClientResponse clientResponse = ClientResponse
                .create(this.getDelegate().getStatusCode(), getMessageReaders())
                .body(Flux.from(body)).build();
        //是否是错误返回
        boolean errorResp = !exchange.getResponse().getStatusCode().is2xxSuccessful() || headers.containsKey(CustomRequestHeaders.HEADER_ORIGIN_HTTP_STATUS);
        Mono<byte[]> bodyMono = clientResponse.bodyToMono(byte[].class).map((bytes) -> {
        	if(isGzip) {
        		bytes = gzipDecode(bytes);
        	}
        	exchange.getAttributes().put(GatewayConstants.CONTEXT_RESPONSE_BODY_SIZE, bytes.length);
        	bodyString = new String(bytes, StandardCharsets.UTF_8);
        	if(debugMode) {
        		logger.info(" origin resp bodyString:{}",StringUtils.left(bodyString, 1024));
        	}
        	//处理response日志
            ApiRequestLog apiLog = exchange.getAttribute(RequestLogCollector.CURRENT_API_LOG_CONTEXT_NAME);
            if(apiLog != null) {
        		//业务异常码
                if(errorResp 
                		&& JsonUtils.isJsonObjectString(bodyString)
                		&& StringUtils.contains(bodyString, GlobalConstants.PARAM_CODE)
                		&& StringUtils.contains(bodyString, GlobalConstants.PARAM_BIZ_CODE)) {
                	String exceptionCode = JsonUtils.getJsonNodeValue(bodyString, GlobalConstants.PARAM_BIZ_CODE);
                	apiLog.setExceptionCode(exceptionCode);
                }
        	}
        	//
            if(subFilterHandlers != null) {
            	String rewriteBody = null;
            	for (SubFilterHandler handler : subFilterHandlers) {
            		rewriteBody = handler.postHandle(exchange,!errorResp, bodyString);
            		if(rewriteBody != null) {
            			bodyString = rewriteBody;
            		}
				}
            }
        	//
        	bytes = bodyString.getBytes();
        	if(isGzip) {
        		bytes = gzipEncode(bodyString.getBytes());
        	}
            return bytes;
        });

        BodyInserter<Mono<byte[]>, ReactiveHttpOutputMessage> bodyInserter = BodyInserters.fromPublisher(bodyMono,byte[].class);
        CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(
                exchange, exchange.getResponse().getHeaders());
        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                .then(Mono.defer(() -> {
                    Flux<DataBuffer> messageBody = outputMessage.getBody();
                    //HttpHeaders headers = getDelegate().getHeaders();
                    if (!withTransferEncodingHeader) {
                        messageBody = messageBody.doOnNext(data -> headers
                                .setContentLength(data.readableByteCount()));
                    }
                    // TODO: fail if isStreamingMediaType?
                    return getDelegate().writeWith(messageBody);
                }));
	}
	
	
	private static byte[] gzipDecode(byte[] encoded) {
		try {
			ByteArrayInputStream bis = new ByteArrayInputStream(encoded);
			GZIPInputStream gis = new GZIPInputStream(bis);
			return FileCopyUtils.copyToByteArray(gis);
		} catch (IOException e) {
			throw new IllegalStateException("couldn't decode body from gzip", e);
		}
	}

	private static byte[] gzipEncode(byte[] original) {
		try {
			ByteArrayOutputStream bis = new ByteArrayOutputStream();
			GZIPOutputStream gos = new GZIPOutputStream(bis);
			FileCopyUtils.copy(original, gos);
			return bis.toByteArray();
		} catch (IOException e) {
			throw new IllegalStateException("couldn't encode body to gzip", e);
		}
	}
}
